#include "kernel_operator.h"
#include <type_traits>
using namespace AscendC;

constexpr int32_t BUFFER_NUM = 2;
template<typename TYPE_BOXES, typename TYPE_SELECTED_INDICES> class KernelNonMaxSuppression {
public:
    __aicore__ inline KernelNonMaxSuppression() {}
    __aicore__ inline void Init(GM_ADDR boxes, GM_ADDR scores, GM_ADDR max_output_boxes_per_class, GM_ADDR iou_threshold, GM_ADDR score_threshold, GM_ADDR selected_indices, uint32_t totalLength, uint32_t ALIGN_NUM, uint32_t block_size, uint32_t core_size, uint32_t core_remain, int32_t num_batches, int32_t num_classes, int32_t spatial_dimension, int center_point_box){
        ASSERT(GetBlockNum() != 0 && "block dim can not be zero!");
        this->blockLength = core_size + (GetBlockNum() == GetBlockIdx() + 1 ? core_remain : 0);   //判断任务配置的核数是否与当前核的索引加一相等
        // this->blockLength = GetBlockNum() == GetBlockIdx() + 1 ? core_remain : core_size; 
        this->tileLength = block_size;  //64个元素，正好16个boxes,一个box4元素.
        this->blockLength = this->blockLength + (this->blockLength % ALIGN_NUM ? ALIGN_NUM - this->blockLength % ALIGN_NUM : 0);
        this->tileNum = this->blockLength / this->tileLength + (this->blockLength % this->tileLength > 0); //1次循环
        this->center_point_box = center_point_box;
        this->num_batches = num_batches;
        this->num_classes = num_classes;
        this->spatial_dimension = spatial_dimension;
        
        auto startPointer = core_size * GetBlockIdx();
        auto bufferlength = this->blockLength;
        //输入参数
        boxesGm.SetGlobalBuffer((__gm__ TYPE_BOXES*)boxes + startPointer, num_batches * spatial_dimension *4);
        scoresGm.SetGlobalBuffer((__gm__ TYPE_BOXES*)scores + startPointer, num_batches * num_classes * spatial_dimension);
        mobpcGm.SetGlobalBuffer((__gm__ TYPE_SELECTED_INDICES*)max_output_boxes_per_class, 1);
        int32_t maxoutputboxes_perclass = mobpcGm.GetValue(0);
        this->maxoutputboxes_perclass = maxoutputboxes_perclass;
        int32_t num_selected_indices = this->num_batches * this->num_classes * this->maxoutputboxes_perclass;
        //输入参数的阈值
        iouthGm.SetGlobalBuffer((__gm__ TYPE_BOXES*)iou_threshold, 1);
        scorethGm.SetGlobalBuffer((__gm__ TYPE_BOXES*)score_threshold, 1);
        //输出参数
        selticsGm.SetGlobalBuffer((__gm__ TYPE_SELECTED_INDICES*)selected_indices + startPointer, num_selected_indices * 3);

        pipe.InitBuffer(boxesInQ, BUFFER_NUM, this->tileLength * sizeof(TYPE_BOXES));
        pipe.InitBuffer(scoresInQ, BUFFER_NUM, this->spatial_dimension * sizeof(TYPE_BOXES));
        pipe.InitBuffer(sortInQ, BUFFER_NUM, this->spatial_dimension *2 * sizeof(TYPE_BOXES));
        pipe.InitBuffer(sresInQ, BUFFER_NUM, this->spatial_dimension * sizeof(TYPE_BOXES));// 排序后的分数队列
        pipe.InitBuffer(iresInQ, BUFFER_NUM, this->spatial_dimension * sizeof(TYPE_BOXES));// 排序后的索引队列
        pipe.InitBuffer(invalInQ, BUFFER_NUM, this->spatial_dimension * sizeof(int32_t));  // float->int32  排序好的框的索引
        pipe.InitBuffer(indexInQ, BUFFER_NUM, this->spatial_dimension * sizeof(int32_t)); // 自己设的升序索引队列
        //pipe.InitBuffer(indexInQ1, BUFFER_NUM, this->spatial_dimension * sizeof(uint32_t));
        //创建储存boxes的值的队列  0：[y1,x1,y2,x2], 1: [x_center,y_center,width,height

        pipe.InitBuffer(selticsOutQ, BUFFER_NUM, this->maxoutputboxes_perclass * 3 * sizeof(TYPE_SELECTED_INDICES));  
        //总的排序好的输出结果   （批次，类别，框)  

    }
    __aicore__ inline void Process() {
        uint32_t length = this->tileLength;
        int32_t spatial_dimension = this->spatial_dimension;
        for(int32_t i = 0; i < num_batches; i++){
            for(int32_t j = 0; j < num_classes; j++){
                CopyIn(i, j, length, spatial_dimension);
                Compute(i, j, length, spatial_dimension);
                CopyOut(i, j, length, spatial_dimension);
            }
        }
    }
private:
    __aicore__ inline void CopyIn(int32_t i, int32_t j, uint32_t length, int32_t spatial_dimension) {
        LocalTensor<TYPE_BOXES> boxes = boxesInQ.AllocTensor<TYPE_BOXES>();
        LocalTensor<TYPE_BOXES> scores = scoresInQ.AllocTensor<TYPE_BOXES>();
        LocalTensor<TYPE_BOXES> sort = sortInQ.AllocTensor<TYPE_BOXES>();
        LocalTensor<TYPE_BOXES> sres = sresInQ.AllocTensor<TYPE_BOXES>();
        LocalTensor<TYPE_BOXES> ires = iresInQ.AllocTensor<TYPE_BOXES>();
        LocalTensor<int32_t> inval = invalInQ.AllocTensor<int32_t>();
        LocalTensor<int32_t> index = indexInQ.AllocTensor<int32_t>();
        // Duplicate(sthval, scorethreshold_value, spatial_dimension);
        DataCopy(boxes, boxesGm[i * this->tileLength], length);
        DataCopy(scores, scoresGm[i * this->num_classes * this-> spatial_dimension + j * this->spatial_dimension], spatial_dimension);
        // Duplicate(index, static_cast<uint32_t>(0.0), spatial_dimension);
        for(uint32_t m = 0; m < spatial_dimension; m++){
            uint32_t n = 0;
            index.SetValue(m, n);
            n++;
        }
        // CreateVecIndex(index, (int32_t)0, this->spatial_dimension);
        boxesInQ.EnQue(boxes);
        scoresInQ.EnQue(scores);
        indexInQ.EnQue(index);
        sortInQ.EnQue(sort);
        sresInQ.EnQue(sres);
        iresInQ.EnQue(ires);
        invalInQ.EnQue(inval);
    }
    __aicore__ inline void Compute(int32_t i, int32_t j, uint32_t length, int32_t spatial_dimension) {
        LocalTensor<TYPE_BOXES> boxes = boxesInQ.DeQue<TYPE_BOXES>();
        LocalTensor<TYPE_BOXES> scores = scoresInQ.DeQue<TYPE_BOXES>();
        LocalTensor<TYPE_BOXES> sort = sortInQ.DeQue<TYPE_BOXES>();
        LocalTensor<TYPE_BOXES> sres = sresInQ.DeQue<TYPE_BOXES>();
        LocalTensor<TYPE_BOXES> ires = iresInQ.DeQue<TYPE_BOXES>();
        LocalTensor<int32_t> inval = invalInQ.DeQue<int32_t>();
        LocalTensor<uint32_t> index = indexInQ.DeQue<uint32_t>();
        LocalTensor<TYPE_SELECTED_INDICES> selectindex = selticsOutQ.AllocTensor<TYPE_SELECTED_INDICES>();
        //boxes的数值队列
        uint32_t mask = 0;
        uint64_t rsvdCnt = 16;
        // auto bits = bitsInQ.Get<uint8_t>();
        // float y1,x1,y2,x2,x_center,y_center,width,height;
        // int center_point_box = this->center_point_box;  
        float iou_threshold_value = iouthGm.GetValue(0);
        float score_threshold_value = scorethGm.GetValue(0);
        // Duplicate(sthval, score_threshold_value, spatial_dimension);
        //第一步将那些分数大于阈值的标记为1，小于阈值的标记为0
        // Compare(bits, scores, sthval, CMPMODE::GE, length);
        //第二步根据间隔的mask,对bits为1的那些进行scores的排序(从大到小)
        //使用循环给index赋值0~15
        Sort32(sort, scores, index, 1);
        GatherMask(sres, sort, static_cast<uint8_t>(3), false, mask, {1,1,0,0}, rsvdCnt);
        GatherMask(ires, sort, static_cast<uint8_t>(4), false, mask, {1,1,0,0}, rsvdCnt);
        Cast(inval, ires, RoundMode::CAST_RINT, spatial_dimension);
        //第三步，根据间隔mask.对分数高的框进行iou计算与iou阈值比较，比阈值高的分数较低的框去除
        //之后再进行排序，取前max_output_boxes_per_class位的框
        // Sort();
        selectindex.SetValue(0, i);  //批次索引
        selectindex.SetValue(1, j);  //类别索引
        selectindex.SetValue(2, inval.GetValue(3));  //框索引
        if(this->center_point_box == 0){           //[y1, x1, y2, x2],areas = (x2-x1)*(y2-y1)
            for(int32_t k = 0; k < spatial_dimension; k++){
                float box1_y1 = sres.GetValue(k * 4);
                float box1_x1 = sres.GetValue(k * 4 + 1);
                float box1_y2 = sres.GetValue(k * 4 + 2);
                float box1_x2 = sres.GetValue(k * 4 + 3);
               
                for(int32_t h = k + 1; h < spatial_dimension; h++){
                    float box2_y1 = sres.GetValue(h * 4);
                    float box2_x1 = sres.GetValue(h * 4 + 1);
                    float box2_y2 = sres.GetValue(h * 4 + 2);
                    float box2_x2 = sres.GetValue(h * 4 + 3);

                    float box1_area = (box1_y2 - box1_y1) * (box1_x2 - box1_x1);
                    float box2_area = (box2_y2 - box2_y1) * (box2_x2 - box2_x1);

                    float y1_inter = max(box1_y1, box2_y1);
                    float x1_inter = max(box1_x1, box2_x1);
                    float y2_inter = min(box1_y2, box2_y2);
                    float x2_inter = min(box1_y2, box2_x2);

                    float intersection_area = max(0.0f, x2_inter - x1_inter) * max(0.0f, y2_inter - y1_inter);
                    float union_area = box1_area + box2_area - intersection_area;
                    float iou_cal = intersection_area / union_area;

                    if(iou_cal < iou_threshold_value){
                        if(sres.GetValue(h) >= score_threshold_value){
                            selectindex.SetValue(3 * h, i);
                            selectindex.SetValue(3 * h + 1, j);
                            selectindex.SetValue(3 * h + 2, inval.GetValue(h));
                        }
                    }
                }
            }
        }
        else if(this->center_point_box == 1){      //[x_center, y_center, width, height]
            for(int32_t k = 0; k < spatial_dimension; k+=2){
                float box1_y1 = boxes.GetValue(k * 4);
                float box1_x1 = boxes.GetValue(k * 4 + 1);
                float box1_y2 = boxes.GetValue(k * 4 + 2);
                float box1_x2 = boxes.GetValue(k * 4 + 3);

                for(int32_t h = k + 1; h < spatial_dimension; h+=2){
                    float box2_y1 = boxes.GetValue(h * 4);
                    float box2_x1 = boxes.GetValue(h * 4 + 1);
                    float box2_y2 = boxes.GetValue(h * 4 + 2);
                    float box2_x2 = boxes.GetValue(h * 4 + 3);

                    float box1_area = (box1_y2 - box1_y1) * (box1_x2 - box1_x1);
                    float box2_area = (box2_y2 - box2_y1) * (box2_x2 - box2_x1);

                    float y1_inter = max(box1_y1, box2_y1);
                    float x1_inter = max(box1_x1, box2_x1);
                    float y2_inter = min(box1_y2, box2_y2);
                    float x2_inter = min(box1_y2, box2_x2);

                    float intersection_area = max(0.0f, x2_inter - x1_inter) * max(0.0f, y2_inter - y1_inter);
                    float union_area = box1_area + box2_area - intersection_area;
                    float iou_cal = intersection_area / union_area;

                    if(iou_cal < iou_threshold_value){
                        if(sres.GetValue(k) >= score_threshold_value){
                            selectindex.SetValue(3 * k, i);
                            selectindex.SetValue(3 * k + 1, j);
                            selectindex.SetValue(3 * k + 2, inval.GetValue(k));
                        }
                        if(sres.GetValue(h) >= score_threshold_value){
                            selectindex.SetValue(3 * h, i);
                            selectindex.SetValue(3 * h + 1, j);
                            selectindex.SetValue(3 * h + 2, inval.GetValue(h));
                        }
                    }
                    else{
                         if(sres.GetValue(k) >= score_threshold_value){
                            selectindex.SetValue(3 * k, i);
                            selectindex.SetValue(3 * k + 1, j);
                            selectindex.SetValue(3 * k + 2, inval.GetValue(k));
                        }
                    }
                }
            }
        }
        selticsOutQ.EnQue(selectindex);
        boxesInQ.FreeTensor(boxes);
        scoresInQ.FreeTensor(scores);
        indexInQ.FreeTensor(index);
        sortInQ.FreeTensor(sort);
        sresInQ.FreeTensor(sres);
        iresInQ.FreeTensor(ires);
        invalInQ.FreeTensor(inval);
    }
    __aicore__ inline void CopyOut(int32_t i, int32_t j, uint32_t length, int32_t spatial_dimension) {
        LocalTensor<TYPE_SELECTED_INDICES> selectindex = selticsOutQ.DeQue<TYPE_SELECTED_INDICES>();
        DataCopy(selticsGm[i * num_classes * 3 * this->maxoutputboxes_perclass + j * 3 * this->maxoutputboxes_perclass], selectindex, this->maxoutputboxes_perclass * 3);
        selticsOutQ.FreeTensor(selectindex);
    }
private:
    TPipe pipe;
    TQue<QuePosition::VECIN, BUFFER_NUM> boxesInQ, scoresInQ, indexInQ, sortInQ, sresInQ, iresInQ, invalInQ;
    TQue<QuePosition::VECOUT, BUFFER_NUM> selticsOutQ;
    // TBuf<QuePosition::VECCALC> bitsInQ;
    GlobalTensor<TYPE_BOXES> boxesGm, scoresGm, iouthGm, scorethGm;
    GlobalTensor<TYPE_SELECTED_INDICES> mobpcGm, selticsGm;
    uint32_t blockLength;
    uint32_t tileNum;
    uint32_t tileLength;
    int center_point_box;
    int32_t num_batches;
    int32_t num_classes;
    int32_t spatial_dimension;
    int32_t num_selected_indices;
    // float iouthreshold_value;
    // float scorethreshold_value;
    int32_t maxoutputboxes_perclass;
};

extern "C" __global__ __aicore__ void non_max_suppression(GM_ADDR boxes, GM_ADDR scores, GM_ADDR max_output_boxes_per_class, GM_ADDR iou_threshold, GM_ADDR score_threshold, GM_ADDR selected_indices, GM_ADDR workspace, GM_ADDR tiling) {
    GET_TILING_DATA(tiling_data, tiling);
    // TODO: user kernel impl
    KernelNonMaxSuppression<DTYPE_BOXES, DTYPE_SELECTED_INDICES> op;
    op.Init(boxes, scores, max_output_boxes_per_class, iou_threshold, score_threshold, selected_indices, 
    tiling_data.totalLength, tiling_data.ALIGN_NUM, tiling_data.block_size, tiling_data.core_size, tiling_data.core_remain, tiling_data.num_batches, tiling_data.num_classes, tiling_data.spatial_dimension, tiling_data.center_point_box);
    op.Process();
}